##########################################################################################

library(data.table)
library(optparse)
library(Mfuzz)
library(parallel)
library(dplyr)
library(limma)
library(RColorBrewer)

##########################################################################################

option_list <- list(
    make_option(c("--sample_list_file"), type = "character"),
    make_option(c("--rsem_file"), type = "character"),
    make_option(c("--diff_file"), type = "character"),
    make_option(c("--gtf_file"), type = "character"),
    make_option(c("--q_t"), type = "character"),
    make_option(c("--foldchange_t"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    sample_list_file <- "~/20220915_gastric_multiple/rna_combine/analysis/config/tumor_normal.list"
    diff_file <- "~/20220915_gastric_multiple/rna_combine/analysis/images/DiffGene/DiffGene.tsv"
    rsem_file <- "~/20220915_gastric_multiple/rna_combine/analysis/images/DiffGene/CombineCounts.FilterLowExpression-MergeMutiSample.TMM.tsv"
    out_path <- "~/20220915_gastric_multiple/rna_combine/analysis/images/mfuzz_v2"
    gtf_file <- "~/ref/GTF/gencode.v19.ensg_genename.txt"
    q_t <- 0.05
    foldchange_t <- 1.5

}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

sample_list_file <- opt$sample_list_file
rsem_file <- opt$rsem_file
diff_file <- opt$diff_file
out_path <- opt$out_path
gtf_file <- opt$gtf_file
q_t <- as.numeric(opt$q_t)
foldchange_t <- as.numeric(opt$foldchange_t)

dir.create(out_path , recursive = T)

##########################################################################################

info <- data.frame(fread(sample_list_file))
dat_tpm <- data.frame(fread(rsem_file))
dat_diff <- data.frame(fread(diff_file))
dat_gtf <- data.frame(fread(gtf_file , header = F))
colnames(dat_gtf) <- c("gene_id" , "Hugo_Symbol")

##########################################################################################

class_type <- c( "Normal" , "IM" , "IGC" , "DGC")

##########################################################################################

dat_tpm_all <- dat_tpm[,c(ncol(dat_tpm) , 1:(ncol(dat_tpm)-1))]

##########################################################################################
## 提取显著差异的基因
#dat_diff <- merge( dat_diff , dat_gtf , by = "gene_id" )
diff_gene <- subset(dat_diff , padj < q_t & abs(log2FoldChange) > log2(foldchange_t) )

## Normal和IM的差异基因
diff_normal_im <- subset( diff_gene , class2=="Normal" & class1=="IM")$gene_id

## Normal和IGC的差异基因
diff_normal_igc <- subset( diff_gene , class2=="Normal" & class1=="IGC")$gene_id

## Normal和DGC的差异基因
diff_normal_dgc <- subset( diff_gene , class2=="Normal" & class1=="DGC")$gene_id

## IM和iGC的差异基因
diff_im_igc <- subset( diff_gene , class2=="IM" & class1=="IGC")$gene_id

## IM和DGC的差异基因
diff_im_dgc <- subset( diff_gene , class2=="IM" & class1=="DGC")$gene_id

##########################################################################################
## 标记每个cluster对应多少基因
mfuzz.plot_revise <- function (eset, cl, mfrow = c(1, 1), colo, min.mem = 0, time.labels, 
    new.window = TRUE){

    clusterindex <- cl[[3]]
    memship <- cl[[4]]
    memship[memship < min.mem] <- -1
    colorindex <- integer(dim(exprs(eset))[[1]])

    if (missing(colo)) {
        colo <- c("#FF8F00", "#FFA700", "#FFBF00", "#FFD700", 
            "#FFEF00", "#F7FF00", "#DFFF00", "#C7FF00", "#AFFF00", 
            "#97FF00", "#80FF00", "#68FF00", "#50FF00", "#38FF00", 
            "#20FF00", "#08FF00", "#00FF10", "#00FF28", "#00FF40", 
            "#00FF58", "#00FF70", "#00FF87", "#00FF9F", "#00FFB7", 
            "#00FFCF", "#00FFE7", "#00FFFF", "#00E7FF", "#00CFFF", 
            "#00B7FF", "#009FFF", "#0087FF", "#0070FF", "#0058FF", 
            "#0040FF", "#0028FF", "#0010FF", "#0800FF", "#2000FF", 
            "#3800FF", "#5000FF", "#6800FF", "#8000FF", "#9700FF", 
            "#AF00FF", "#C700FF", "#DF00FF", "#F700FF", "#FF00EF", 
            "#FF00D7", "#FF00BF", "#FF00A7", "#FF008F", "#FF0078", 
            "#FF0060", "#FF0048", "#FF0030", "#FF0018")
    }
    colorseq <- seq(0, 1, length = length(colo))
    for (j in 1:max(clusterindex)) {
        tmp <- exprs(eset)[clusterindex == j, , drop = FALSE]
        tmpmem <- memship[clusterindex == j, j]
        if (((j - 1)%%(mfrow[1] * mfrow[2])) == 0) {
            if (new.window) 
                X11()
            par(mfrow = mfrow)
            if (sum(clusterindex == j) == 0) {
                ymin <- -1
                ymax <- +1
            }
            else {
                ymin <- min(tmp)
                ymax <- max(tmp)
            }
            plot.default(x = NA, xlim = c(1, dim(exprs(eset))[[2]]), 
                ylim = c(ymin, ymax), xlab = "Time", ylab = "Expression changes", 
                main = paste("Cluster", j , "(" , nrow(tmp) , ")"), axes = FALSE)
            if (missing(time.labels)) {
                axis(1, 1:dim(exprs(eset))[[2]], c(1:dim(exprs(eset))[[2]]))
                axis(2)
            }
            else {
                axis(1, 1:dim(exprs(eset))[[2]], time.labels)
                axis(2)
            }
        }
        else {
            if (sum(clusterindex == j) == 0) {
                ymin <- -1
                ymax <- +1
            }
            else {
                ymin <- min(tmp)
                ymax <- max(tmp)
            }
            plot.default(x = NA, xlim = c(1, dim(exprs(eset))[[2]]), 
                ylim = c(ymin, ymax), xlab = "Time", ylab = "Expression changes", 
                main = paste("Cluster", j , "(" , nrow(tmp) , ")"), axes = FALSE)
            if (missing(time.labels)) {
                axis(1, 1:dim(exprs(eset))[[2]], c(1:dim(exprs(eset))[[2]]))
                axis(2)
            }
            else {
                axis(1, 1:dim(exprs(eset))[[2]], time.labels)
                axis(2)
            }
        }
        if (!(sum(clusterindex == j) == 0)) {
            for (jj in 1:(length(colorseq) - 1)) {
                tmpcol <- (tmpmem >= colorseq[jj] & tmpmem <= 
                  colorseq[jj + 1])
                if (sum(tmpcol) > 0) {
                  tmpind <- which(tmpcol)
                  for (k in 1:length(tmpind)) {
                    lines(tmp[tmpind[k], ], col = colo[jj])
                  }
                }
            }
        }
    }
}
##########################################################################################
## 按时间表达聚类
plotMfuzz <- function( dat_tpm_all = dat_tpm_all , col_names = col_names , out_path = out_path , tumor = tumor , clust = clust ){

    useClass <- c("Normal" , "IM" , tumor)

    dat_tpm_use <- dat_tpm_all[,col_names]
    DEGs_exp <- dat_tpm_use[dat_tpm_use$gene_id %in% diff_list,]
    time <- data.frame( 
        sample = colnames(DEGs_exp)[-1] ,
        class = sapply( strsplit(colnames(DEGs_exp)[-1] , "_" ) , "[" , 2 ) 
    )

    rownames(DEGs_exp) <- DEGs_exp$gene_id
    DEGs_exp <- DEGs_exp[,-1]

    tmp <- data.frame(colnames(DEGs_exp),t(DEGs_exp))
    temp <- data.frame(time[match(tmp[,1],time[,1]),],tmp)
    DEGs_exp_averp <- t(limma::avereps(temp[,-c(1:3)],ID=temp[,2]))
    DEGs_exp_averp <- DEGs_exp_averp[,col_order]

    ## 1. 预处理：去除表达量太低或者在不同时间点间变化太小的基因等步骤
    # Mfuzz聚类时要求是一个ExpressionSet类型的对象，所以需要先用表达量构建这样一个对象。
    eset <- new("ExpressionSet",exprs = DEGs_exp_averp)

    # 根据标准差去除样本间差异太小的基因
    eset <- filter.std(eset,min.std=0)

    ## 2. 标准化：聚类时需要用一个数值来表征不同基因间的距离，Mfuzz中采用的是欧式距离，
    # 由于普通欧式距离的定义没有考虑不同维度间量纲的不同，所以需要先进行标准化
    eset <- standardise(eset)

    ## 3. 聚类：Mfuzz中的聚类算法需要提供两个参数，第一个参数为希望最终得到的聚类的个数，这个参数由我们直接指定；
    # 第二个参数称之为fuzzifier值，用小写字母m表示，可以通过函数评估一个最佳取值
    #clust <- 9 # 聚类个数，文章中用的6个聚类，我们也用6个
    m <- mestimate(eset) #  评估出最佳的m值
    cl <- mfuzz(eset, c = clust, m = m) # 聚类

    out_name <- paste0(out_path , "/mfuzz_plot_",tumor,".pdf")
    color.2 <- colorRampPalette(rev(c("#ff0000", "Yellow", "OliveDrab1")))(1000)
    pdf(out_name)
    mfuzz.plot_revise(eset,cl,mfrow=c(2,3),new.window= FALSE,time.labels=colnames(DEGs_exp_averp),colo = color.2)
    dev.off()

    # 在cl这个对象中就保存了聚类的完整结果，对于这个对象的常见操作如下
    # cl$size # 查看每个cluster中的基因个数
    # [1] 2269 1982 2289 1653 1393 1729

    # cl$cluster[cl$cluster == 1] # 提取某个cluster下的基因
    # cl$membership # 查看基因和cluster之间的membership

    ## acore函数帮助我们计算出来了每个趋势分组里面的每个基因的重要性
    acore <- acore(eset , cl , min.acore=0)
    acore_list <- do.call(rbind, lapply(seq_along(acore), function(i){ data.frame(CLUSTER=i, acore[[i]])}))
    acore_list <- merge( dat_gtf , acore_list , by.x = "gene_id" , by.y = "NAME" )

    result_acore <- merge( dat_diff , acore_list , by = "gene_id" )
    result_acore <- subset( result_acore , !(class1 %in% class_type[!(class_type %in% useClass)]) & !(class2 %in% class_type[!(class_type %in% useClass)])  )

    out_name <- paste0(out_path , "/mfuzz_plot_",tumor,".tsv")
    write.table( result_acore , out_name , row.names = F , quote = F , sep = "\t" )

}

clust <- 6

## IGC
diff_normal_im_igc <- unique(c(diff_normal_im , diff_normal_igc , diff_im_igc))
diff_list <- diff_normal_im_igc

col_names <- grep( "gene_id|Normal|IM|IGC" , colnames(dat_tpm_all) , value = T )
col_order <- c("Normal" , "IM" , "IGC")
tumor <- "IGC"
set.seed(1234)
plotMfuzz( dat_tpm_all = dat_tpm_all , col_names = col_names , out_path = out_path , tumor = tumor , clust = clust )

## DGC
diff_normal_im_dgc <- unique(c(diff_normal_im , diff_normal_dgc , diff_im_dgc))
diff_list <- diff_normal_im_dgc

col_names <- grep( "gene_id|Normal|IM|DGC" , colnames(dat_tpm_all) , value = T )
col_order <- c("Normal" , "IM" , "DGC")
tumor <- "DGC"
set.seed(1234)
plotMfuzz( dat_tpm_all = dat_tpm_all , col_names = col_names , out_path = out_path , tumor = tumor , clust = clust )
